import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import classification_report, accuracy_score
from joblib import dump

# 规范类名
class Classification:
    def __init__(self):
        # 加载数据
        self.df = pd.read_csv('fina_indicator.csv')
        # 关键：打印CSV中所有实际列名，方便你核对
        print("="*50)
        print("CSV文件中实际存在的列名：")
        for i, col in enumerate(self.df.columns.tolist(), 1):
            print(f"{i:2d}. {col}")
        print("="*50)

    def get_conditions(self):
        """分类前准备（修复列名问题）"""
        df = self.df
        # 计算收益和风险比例
        df['max_ratio'] = df['max_close'] / df['the_close']  # 收益潜力
        df['min_ratio'] = df['min_close'] / df['the_close']  # 风险潜力

        # 自动划分高低阈值
        high_return_threshold = df['max_ratio'].quantile(0.4)  # 前40%定义为高收益
        high_risk_threshold = df['min_ratio'].quantile(0.4)    # 前40%定义为高风险

        # 生成分类标签
        conditions = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold)
        ]
        choices = ['高收益高风险', '低收益高风险', '高收益低风险', '低收益低风险']
        df['category'] = np.select(conditions, choices, default='未知')
        
        # 特征选择：删除不存在的 'tcfcf' 列（核心修复！）
        # 如果你CSV中有类似现金流的列（比如 'tcff'、'total_cf' 等），请替换下面的列名
        features = df[['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 
                       'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta']]

        # 标签编码
        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])

        # 数据标准化
        scaler = StandardScaler()
        X = scaler.fit_transform(features)
        y = df['category_encoded']
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
        dump(scaler, 'scaler.joblib')
        return X_train, X_test, y_train, y_test, label_encoder

    def knn_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """KNN分类器"""
        knn = KNeighborsClassifier(n_neighbors=3)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        print("\n" + "="*50)
        print("KNN分类器评估报告：")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        print(f"KNN准确率：{accuracy_score(y_test, y_pred):.4f}")
        dump(knn, 'knn_model.joblib')
        dump(label_encoder, 'label_encoder.joblib')

    def svc_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """SVC分类器"""
        svc = SVC(kernel='linear', C=1.0, random_state=42)
        svc.fit(X_train, y_train)
        y_pred = svc.predict(X_test)
        print("\n" + "="*50)
        print("SVC分类器评估报告：")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        print(f"SVC准确率：{accuracy_score(y_test, y_pred):.4f}")
        dump(svc, 'svc_model.joblib')

    def decision_tree_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        """决策树分类器"""
        dt = DecisionTreeClassifier(random_state=42, max_depth=8)  # 限制树深避免过拟合
        dt.fit(X_train, y_train)
        y_pred = dt.predict(X_test)
        print("\n" + "="*50)
        print("决策树分类器评估报告：")
        print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))
        print(f"决策树准确率：{accuracy_score(y_test, y_pred):.4f}")
        dump(dt, 'dt_model.joblib')

if __name__ == '__main__':
    # 实例化类（自动打印列名）
    ci = Classification()
    # 准备数据（此时已无 tcfcf 列，不会报错）
    X_train, X_test, y_train, y_test, label_encoder = ci.get_conditions()
    # 训练所有分类器
    ci.knn_utils(X_train, X_test, y_train, y_test, label_encoder)
    ci.svc_utils(X_train, X_test, y_train, y_test, label_encoder)
    ci.decision_tree_utils(X_train, X_test, y_train, y_test, label_encoder)
    
    print("\n" + "="*50)
    print("所有模型训练完成！已保存：scaler.joblib、knn_model.joblib、svc_model.joblib、dt_model.joblib")
